import numpy as np
from scipy.ndimage.interpolation import zoom

def down_sample(data, factor:int=2):
    image, label = data
    assert factor>=2
    assert image.shape[:2] == label.shape[:2]
    if len(image.shape)==3:
        image = zoom(image, (1/factor, 1/factor, 1), order=3)
    else:
        image = zoom(label, (1/factor, 1/factor), order=3)
    label = zoom(label, (1/factor, 1/factor), order=0)

    return image, label


def make_square(image, pad_value=0):
    # Get image dimensions
    height, width = image.shape[:2]
    
    # Identify the smaller dimension
    if height > width:
        diff = height - width
        padding = ((0, 0), (diff // 2, diff - diff // 2))
    else:
        diff = width - height
        padding = ((diff // 2, diff - diff // 2), (0, 0))
        
    # If the image is grayscale
    if len(image.shape) == 2:
        return np.pad(image, padding, mode='constant', constant_values=pad_value)
    
    # If the image is colored
    padding = padding + ((0, 0),)
    return np.pad(image, padding, mode='constant', constant_values=pad_value)
